package com.zengde.ai_sandbox.fluxtest;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.disposables.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

public class FluxWrapperExample {

    public static void main(String[] args) {
        String apiKey = System.getenv("ARK_API_KEY");
        ArkService service = ArkService.builder().apiKey(apiKey).build();

        // 创建请求消息
        final List<ChatMessage> streamMessages = new ArrayList<>();
        streamMessages.add(ChatMessage.builder().role(ChatMessageRole.SYSTEM)
                .content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build());
        streamMessages.add(ChatMessage.builder().role(ChatMessageRole.USER)
                .content("常见的十字花科植物有哪些？").build());

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model("doubao-1-5-pro-32k-250115")
                .messages(streamMessages)
                .build();

        // 创建包装后的 Flux 流
        Flux<String> wrappedFlux = wrapDoubaoStream(service, request);

        // 订阅并处理转换后的流
        wrappedFlux.subscribe(
                transformedContent -> {
                    System.out.println("Received transformed: " + transformedContent);
                    // 这里可以添加自定义处理逻辑
                },
                error -> System.err.println("Error occurred: " + error.getMessage()),
                () -> {
                    System.out.println("Stream completed");
                    service.shutdownExecutor();
                }
        );
    }

    /**
     * 包装豆包流式响应为 Flux，并转换数据内容
     */
    private static Flux<String> wrapDoubaoStream(ArkService service, ChatCompletionRequest request) {
        return Flux.create(emitter -> {
            // 创建原子引用以支持取消操作
            AtomicReference<Disposable> disposableRef = new AtomicReference<>();

            // 启动豆包流式请求
            Disposable disposable = service.streamChatCompletion(request)
                    .doOnError(emitter::error)  // 错误传递
                    .doOnCancel(() -> {
                        if (!emitter.isCancelled()) {
                            emitter.complete();
                        }
                    })
                    .subscribe(
                            chunk -> {
                                if (chunk.getChoices() != null && !chunk.getChoices().isEmpty()) {
                                    // 获取原始内容
                                    String originalContent = (String) chunk.getChoices().get(0).getMessage().getContent();

                                    // 转换内容（示例：添加前缀）
                                    String transformed = transformContent(originalContent);

                                    // 发射转换后的内容
                                    emitter.next(transformed);
                                }
                            },
                            emitter::error,  // 错误传递
                            emitter::complete  // 完成信号传递
                    );

            // 注册取消回调
            disposableRef.set(disposable);
            emitter.onCancel(() -> {
                Disposable d = disposableRef.get();
                if (d != null && !d.isDisposed()) {
                    d.dispose();
                }
            });
        }, FluxSink.OverflowStrategy.BUFFER);  // 设置背压策略
    }

    /**
     * 内容转换函数（示例：添加前缀）
     */
    private static String transformContent(String original) {
        // 这里可以实现任意转换逻辑
        return "[TRANSFORMED] " + original;
    }
}